#include "PoissonSampler.h"
#include "ompl/base/SpaceInformation.h"
#include "ompl/base/spaces/RealVectorBounds.h"
#include "ompl/base/spaces/SE2StateSpace.h"
#include <boost/math/constants/constants.hpp>
#include "PlanningParams.h"
#include <time.h>

using namespace ompl;

namespace poisson
{

void printState(const base::SpaceInformation *si, const base::State* state, const char* stateName)
{
	std::vector<double> vData(7);
	si->getStateSpace()->serialize(&(vData[0]), state);
	printf("State %s : ", stateName);
	for (int i = 0; i < 7; ++i)
	{
		printf("%7.3f ", vData[i]);
	}
	printf("\n");
}

////////////////////////////////////////////////////////////////////////////////

// returns an instance
base::ValidStateSamplerPtr PoissonSampler::allocPoissonStateSampler(const base::SpaceInformation *si)
{
	PoissonSampler* p = new PoissonSampler(si);
	return base::ValidStateSamplerPtr(p);
}

////////////////////////////////////////////////////////////////////////////////

PoissonSampler::PoissonSampler(const base::SpaceInformation *si) :
ValidStateSampler(si), sampler_(si->allocStateSampler()), initialized_(false), resetRequired_(false),
uiDimension_(0), dBaseRadius_((double)BASE_RADIUS), uiNumTotalCells_(0), startState_(NULL), goalState_(NULL)
{
	name_ = "poisson";
}

bool PoissonSampler::sample(ompl::base::State *state)
{
	if (!initialized_)
		setup();
	if (resetRequired_)
		reset();

	PoissonState* poissonState = dynamic_cast<PoissonState*>(state);

	int uiSampleId, uiSampleLevel;

	// TODO?

	return false;
}

bool PoissonSampler::sampleNear(ompl::base::State *state, const ompl::base::State *near, const double distance)
{
	return false;
}

PoissonState* PoissonSampler::addStartPosition(const ompl::base::State* state)
{
	startState_ = addState(-1, 0, state, false);

	// start position is in the tree
	startState_->setTreeNode(true);

	return startState_;
}

PoissonState* PoissonSampler::addGoalPosition(const ompl::base::State* state)
{
	goalState_ = addState(-1, 0, state, true);

	return goalState_;
}

void PoissonSampler::setup()
{
	initializeVariablesForTemplate();
	createAdaptiveTemplate();

	// precompute samples
	initializeVariables();
	precomputeSamples();

	initialized_ = true;
}

void PoissonSampler::reset()
{
	resetRequired_ = false;
}

////////////////////////////////////////////////////////////////////////////////

void PoissonSampler::initializeVariables()
{
	uiDimension_ = si_->getStateDimension();

	getSpaceBounds();

	vCellWidth_.resize(MAX_LEVEL);
	
	vCellWidth_[0] = dBaseRadius_ / (2.0 * sqrt(3.0));
	for (int i = 1; i < MAX_LEVEL; ++i)
	{
		vCellWidth_[i] = vCellWidth_[i - 1] * 0.5f;
	}
	vRadius_.resize(MAX_LEVEL);
	vRadius_[0] = dBaseRadius_;
	for (int i = 1; i < MAX_LEVEL; ++i)
	{
		vRadius_[i] = vRadius_[i - 1] * 0.5f;
	}

	vvDimensionSize_.resize(MAX_LEVEL);
	vvAccumulatedDimensionSize_.resize(MAX_LEVEL);
	for (int k = 0; k < MAX_LEVEL; ++k)
	{
		vvDimensionSize_[k].resize(uiDimension_);
		vvAccumulatedDimensionSize_[k].resize(uiDimension_);
		for (unsigned int i = 0; i < uiDimension_; ++i)
		{
			double dRange = vBoundsMax_[i] - vBoundsMin_[i];
			vvDimensionSize_[k][i] = (int)(ceil(dRange / vCellWidth_[k]));
			vvAccumulatedDimensionSize_[k][i] = vvDimensionSize_[k][i];

			for (unsigned int j = 0; j < i; ++j)
			{
				vvAccumulatedDimensionSize_[k][i] *= vvDimensionSize_[k][j];
			}
		}
	}
	uiNumTotalCells_ = vvAccumulatedDimensionSize_[0][uiDimension_ - 1];

	if (vBaseCells_.size() < vvAccumulatedDimensionSize_[0][uiDimension_ - 1])
		vBaseCells_.resize(vvAccumulatedDimensionSize_[0][uiDimension_ - 1]);
	vvAdditionalCells_.resize(MAX_LEVEL);

	// 2 ^ d
	int iNumSubcells = (1 << uiDimension_);

	// max : 9^6 = 531441
	vSubcellIds_.reserve(iNumSubcells * 10000);

	// TODO: should be enough (or pointers will be invalid)
	vSamples_.reserve(100000);

	// allocate two active cell vectors for dart throwing
	vActiveCells_.reserve(POISSON_DISK_SAMPLING_CONSTANT_B * uiNumTotalCells_);
	vActiveCells2_.reserve(POISSON_DISK_SAMPLING_CONSTANT_B * uiNumTotalCells_);

	nn_.reset(new ompl::NearestNeighborsGNAT<PoissonState*>());
	nn_->setDistanceFunction(boost::bind(&PoissonSampler::distanceFunctionPoissonState, this, _1, _2));
}

void PoissonSampler::precomputeSamples()
{
	clock_t timeStart = clock();

	if (vBaseCells_.size() < vvAccumulatedDimensionSize_[0][uiDimension_ - 1])
		vBaseCells_.resize(vvAccumulatedDimensionSize_[0][uiDimension_ - 1]);
	for (int j = 0; j < vvAccumulatedDimensionSize_[0][uiDimension_ - 1]; ++j)
	{
		vBaseCells_[j].baseId = j;
		vBaseCells_[j].level = 0;
		vBaseCells_[j].sampleId = -1;
		idToIndices(j, vBaseCells_[j].indices, 0);
	}

	// generate samples of level 0
	generateSamples();

	clock_t precomputationTime = clock() - timeStart;
	OMPL_INFORM("Time for precomputation : %f", precomputationTime / (double)CLOCKS_PER_SEC);
}

void PoissonSampler::generateSamples()
{
	vActiveCells_.resize(uiNumTotalCells_);
	for (int i = 0; i < uiNumTotalCells_; ++i)
		vActiveCells_[i] = vBaseCells_[i];

	std::vector<GridCell>* pCurrentActiveCells = &vActiveCells_;
	std::vector<GridCell>* pNextActiveCells = &vActiveCells2_;

	throwDarts(pCurrentActiveCells, pNextActiveCells);

	OMPL_INFORM("Poisson-disk samples generated.");
}

// TODO:
class RandIntClass
{
public:
	RandIntClass() {}
	int operator() (int aRange)
	{
		srand(0);
		int result = rand() % aRange;
		return result;
	}
};

void PoissonSampler::throwDarts(std::vector<GridCell>* pCurrentActiveCells, std::vector<GridCell>* pNextActiveCells)
{
	RandIntClass Rand;

	int uiCellLevel = 0;
	int uiNumSamples = vSamples_.size();
	unsigned int uiTrials = __max((unsigned int)ceil(POISSON_DISK_SAMPLING_CONSTANT_A), 1);
	std::vector<bool> vIsCellRemoved;

	base::State* state = si_->allocState();

	// throw darts
	while (pCurrentActiveCells->size() > 0)
	{
		unsigned int uiNumDart = (unsigned int)ceil(POISSON_DISK_SAMPLING_CONSTANT_A * pCurrentActiveCells->size());
		int uiLastNumSamples = uiNumSamples;
		
		std::random_shuffle(pCurrentActiveCells->begin(), pCurrentActiveCells->end(), Rand);

		if (uiCellLevel == MAX_PRECOMPUTATION_LEVEL - 1)
			uiNumDart = pCurrentActiveCells->size();

		double dCellWidth = vCellWidth_[uiCellLevel];

		vIsCellRemoved.clear();
		vIsCellRemoved.reserve(pCurrentActiveCells->size());	
		for (unsigned int i = 0; i < pCurrentActiveCells->size(); ++i)
		{
			GridCell& currentCell = (*pCurrentActiveCells)[i];
			int uiCellId = currentCell.baseId;

			GridCell* baseCell = getGridCell(uiCellId, 0);
			if (baseCell == NULL)
				continue;

			// if parent base grid cell has a sample, just remove the current cell
			if (baseCell->GetStatus() == CS_SAMPLE)
			{
				vIsCellRemoved[i] = true;
				continue;
			}
			if (i >= uiNumDart)
				continue;

			// throw candidate dart
			for (unsigned int j = 0; j < uiTrials; ++j)
			{
				randomPointInCell(currentCell, state, dCellWidth);
				
				// disk-free check
				if (isDiskFree(state, currentCell, 0))
				{
					addState(uiCellId, 0, state, true);

					baseCell->sampleId = uiNumSamples;
					baseCell->level = 0;
					// TODO:
					//sampleManager_.AddSample(uiNumSamples, 0);
					++uiNumSamples;

					vIsCellRemoved[i] = true;
					break;
				}
				else
				{
					vIsCellRemoved[i] = false;
				}
				
			}
		}
		//OMPL_INFORM("[%d] %d/%d (%f)samples added" , uiCellLevel, (uiNumSamples - uiLastNumSamples), uiNumDart, (double)(uiNumSamples - uiLastNumSamples) / uiNumDart);

		// for all active cells, 
		// if level < bit limit subdivide cell
		// retain uncovered subcells to next active set
		if (uiCellLevel < MAX_PRECOMPUTATION_LEVEL - 1)
			++uiCellLevel;
		pNextActiveCells->clear();

		if (uiCellLevel < MAX_PRECOMPUTATION_LEVEL - 1)
		{
			for (unsigned int i = 0; i < pCurrentActiveCells->size(); ++i)
			{
				if (vIsCellRemoved[i] == false)
				{
					GridCell& currentCell = (*pCurrentActiveCells)[i];
					addUncoveredSubcells(currentCell, pNextActiveCells, uiCellLevel, 0);
				}
			}
		}
		//OMPL_INFORM("Active Cell %d size : %d", uiCellLevel, pNextActiveCells->size());

		std::vector<GridCell>* pTemp = pCurrentActiveCells;
		pCurrentActiveCells = pNextActiveCells;
		pNextActiveCells = pTemp;
	}

	si_->freeState(state);
}

GridCell* PoissonSampler::getGridCell(int id, int uiLevel, bool create)
{
	if (!isValidId(id, uiLevel))
		return NULL;

	if (uiLevel == 0)
		return &vBaseCells_[id];
	else
	{
		std::map<int, GridCell>::iterator it = vvAdditionalCells_[uiLevel].find(id);
		if (it != vvAdditionalCells_[uiLevel].end())
			return &(it->second);
		if (create)
		{
			std::vector<int> indices;
			idToIndices(id, indices, uiLevel);
			vvAdditionalCells_[uiLevel][id] = GridCell(id, uiLevel, indices);
			return &(vvAdditionalCells_[uiLevel][id]);
		}
	}
	return NULL;
}

const GridCell* PoissonSampler::getGridCell(int id, int uiLevel) const
{
	if (!isValidId(id, uiLevel))
		return NULL;

	if (uiLevel == 0)
		return &vBaseCells_[id];
	else
	{
		std::map<int, GridCell>::const_iterator it = vvAdditionalCells_[uiLevel].find(id);
		if (it != vvAdditionalCells_[uiLevel].end())
			return &(it->second);
	}
	return NULL;
}

__inline bool PoissonSampler::isValidId(int id, int uiLevel) const
{
	return (id < vvAccumulatedDimensionSize_[uiLevel][uiDimension_ - 1]);
}

void PoissonSampler::randomPointInCell(const GridCell& cell, base::State* state, double dWidth)
{
	std::vector<double> vPoint(si_->getStateSpace()->getSerializationLength() / sizeof(double));
	unsigned int uiDimTranslation = uiDimension_;
	unsigned int uiDimRotation = 0;
	if (si_->getStateSpace()->isCompound())
	{
		uiDimTranslation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->getDimension();
		uiDimRotation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->getDimension();
	}
	
	for (int j = 0; j < uiDimTranslation; ++j)
	{
		double dCellMin = vBoundsMin_[j] + cell.indices[j] * dWidth;
		if (dCellMin + dWidth > vBoundsMax_[j])
			dWidth = vBoundsMax_[j] - dCellMin;
		vPoint[j] = dCellMin + rng_.uniformReal(0, dWidth);
	}

	switch (uiDimRotation)
	{
	case 0:
		break;
	case 1:
		vPoint[uiDimTranslation] = rng_.uniformReal(-boost::math::constants::pi<double>(), boost::math::constants::pi<double>());
		break;
	case 3:
		rng_.quaternion(&vPoint[uiDimTranslation]);
		break;
	default:
		assert(false);
		break;
	}

	si_->getStateSpace()->deserialize(state, &(vPoint[0]));
}

bool PoissonSampler::isDiskFree(const base::State* state, const GridCell& cell, int uiSampleLevel, bool inTemplate) const
{
	if (inTemplate)
	{
		const base::State* nearest = nnTemplate_->nearest(state);
		if (nearest && distanceFunction(nearest, state) < vRadius_[uiSampleLevel])
			return false;

		return true;
	}
	else
	{
		if (nn_->size() == 0)
			return true;

		PoissonState poissonState(si_, state, -1, -1);

		const PoissonState* nearest = nn_->nearest(&poissonState);
		if (nearest && distanceFunction(nearest->getState(), state) < vRadius_[uiSampleLevel])
			return false;

		return true;

	}
}

PoissonState* PoissonSampler::addState(int cid, int lv, const base::State* state, bool addToNN)
{
	PoissonState* newPoissonState = new PoissonState(si_, state, vSamples_.size(), lv);
	vSamples_.push_back(newPoissonState);

	if (addToNN)
		nn_->add(newPoissonState);

	return newPoissonState;
}

void PoissonSampler::getSpaceBounds()
{
	vBoundsMin_.resize(uiDimension_);
	vBoundsMax_.resize(uiDimension_);

	std::vector<base::StateSpacePtr> vSubSpaces = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspaces();
	int uiDimIndex = 0;
	for (unsigned int i = 0; i < vSubSpaces.size(); ++i)
	{
		int uiDim = vSubSpaces[i]->getDimension();
		switch(vSubSpaces[i]->getType())
		{
		case base::STATE_SPACE_REAL_VECTOR:
			{
				const base::RealVectorBounds& bounds = vSubSpaces[i]->as<base::RealVectorStateSpace>()->getBounds();
				for (int j = 0; j < uiDim; ++j)
				{
					vBoundsMin_[uiDimIndex + j] = bounds.low[j];
					vBoundsMax_[uiDimIndex + j] = bounds.high[j];
				}
			}
			break;

		case base::STATE_SPACE_SO2:
			{
				vBoundsMin_[uiDimIndex] = -boost::math::constants::pi<double>();
				vBoundsMax_[uiDimIndex] = boost::math::constants::pi<double>();
			}
			break;

		case base::STATE_SPACE_SO3:
			{
				for (int j = 0; j < uiDim; ++j)
				{
					vBoundsMin_[uiDimIndex + j] = -1 * SO3_SCALE;
					vBoundsMax_[uiDimIndex + j] = 1 * SO3_SCALE;
				}
			}
			break;

		default:
			OMPL_INFORM("Unsupported state space.");
			break;
		}
		uiDimIndex += uiDim;
	}
}

void PoissonSampler::addUncoveredSubcells(GridCell& cell, std::vector<GridCell>* pNextActiveCells, int uiNextLevel, int uiSampleLevel) const
{
	// check subcell vertices are disk-free or not
	int uiCellId = cell.baseId;
	double dCellWidth = vCellWidth_[cell.level];

	// vertices coordinates
	std::vector<std::vector<double> > vCoords;
	vCoords.resize(uiDimension_);
	for (int i = 0; i < uiDimension_; ++i)
	{
		vCoords[i].resize(3);
		vCoords[i][0] = vBoundsMin_[i] + cell.indices[i] * dCellWidth;
		vCoords[i][1] = vCoords[i][0] + 0.5 * dCellWidth;
		vCoords[i][2] = vCoords[i][0] + dCellWidth;
	}
	int numSubcells = (1 << uiDimension_);

	std::vector<bool> bCovered(numSubcells);
	for (int j = 0; j < numSubcells; ++j)
		bCovered[j] = false;

	for (int i = 0; i < uiDimension_; ++i)
	{
		if (vCoords[i][1] > vBoundsMax_[i])
		{
			vCoords[i][1] = vBoundsMax_[i];

			for (int j = 0; j < numSubcells; ++j)
			{
				int mask = (1 << i);
				if (mask & j)
				{
					bCovered[j] = true;
				}
			}
		}
		else if (vCoords[i][2] > vBoundsMax_[i])
			vCoords[i][2] = vBoundsMax_[i];
	}

	int numVertices = 1;
	for (int i = 0; i < uiDimension_ ; ++i)
		numVertices *= 3;

	std::vector<std::vector<double> > vertices(numVertices);
	std::vector<base::State*> vertexStates(numVertices);
	
	for (int i = 0; i < numVertices; ++i)
	{
		vertices[i].resize(si_->getStateSpace()->getSerializationLength() / sizeof(double));
		int div = 1;
		for (int j = 0; j < uiDimension_ ; ++j)
		{
			vertices[i][j] = vCoords[j][(i / div) % 3];
			div *= 3;
		}
		vertexStates[i] = si_->allocState();
		si_->getStateSpace()->deserialize(vertexStates[i], &(vertices[i][0]));
	}

	std::vector<std::vector<double> > halfPoints;
	halfPoints.resize(uiDimension_);
	for (int j = 0; j < uiDimension_ ; ++j)
	{
		halfPoints[j].resize(2);
		halfPoints[j][0] = (vCoords[j][0] + vCoords[j][1]) * 0.5;
		halfPoints[j][1] = (vCoords[j][1] + vCoords[j][2]) * 0.5;
	}

	std::vector<int> vBaseIndices;
	idToIndices(uiCellId, vBaseIndices, uiSampleLevel);

	std::vector<double> vStateCoord(uiDimension_);
	for (int j = 0; j < uiDimension_; ++j)
		vStateCoord[j] = vCoords[j][1];
	base::State* center = si_->allocState();
	si_->getStateSpace()->deserialize(center, &(vStateCoord[0]));
	
	PoissonState poissonCenterState(si_, center, -1, -1);
	std::vector<PoissonState*> neighbors;
	nn_->nearestR(&poissonCenterState, vRadius_[uiSampleLevel], neighbors);
	for (unsigned int i = 0 ; i < neighbors.size(); ++i)
	{
		si_->getStateSpace()->serialize(&(vStateCoord[0]), neighbors[i]->getState());

		for (int k = 0; k < numSubcells; ++k)
		{
			if (bCovered[k])
				continue;
				
			std::vector<int> indices(uiDimension_);
			for (int l = 0; l < uiDimension_; ++l)
			{
				indices[l] = (k >> l) % 2;
			}
				
			std::vector<int> vertexIndices(uiDimension_);
			for (int l = 0; l < uiDimension_; ++l)
			{
				vertexIndices[l] = indices[l] + (vStateCoord[l] < halfPoints[l][indices[l]]) ? 1 : 0;
			}

			int vertexIndex = vertexIndices[0];
			int acc = 1;
			for (int l = 1; l < uiDimension_; ++l)
			{
				acc *= 3;
				vertexIndex += acc * vertexIndices[l];
			}

			if (distanceFunction(neighbors[i], vertexStates[vertexIndex]) < vRadius_[uiSampleLevel])
				bCovered[k] = true;
		}
	}

	si_->freeState(center);
	for (int i = 0; i < numVertices; ++i)
	{
		si_->freeState(vertexStates[i]);
	}

	// add uncovered cell to next active cell vector
	for (int j = 0; j < numSubcells; ++j)
	{
		if (bCovered[j] == false)
		{
			std::vector<int> indices(uiDimension_);
			for (int l = 0; l < uiDimension_; ++l)
			{
				indices[l] = (j >> l) % 2;
			}

			GridCell subCell = GridCell(cell.baseId, uiNextLevel, cell.indices);
			subCell.sampleId = -1;
			for (int l = 0; l < uiDimension_; ++l)
			{
				subCell.indices[l] = cell.indices[l] * 2 + indices[l];
			}
			pNextActiveCells->push_back(subCell);
		}
	}
}

// template functions
////////////////////////////////////////////////////////////////////////////////
void PoissonSampler::initializeVariablesForTemplate()
{
	OMPL_INFORM("Creating adaptive template");

	uiDimension_ = si_->getStateDimension();
	unsigned int uiDimTranslation = uiDimension_;
	unsigned int uiDimRotation = 0;
	if (si_->getStateSpace()->isCompound())
	{
		uiDimTranslation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->getDimension();
		uiDimRotation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->getDimension();
	}

	vCellWidth_.resize(MAX_LEVEL);
	vCellWidth_[0] = dBaseRadius_ / (2.0 * sqrt(3.0));

	int maxNeighborCellInRow = (int)ceil(1.5 * dBaseRadius_ / vCellWidth_[0]) * 2 + 1;

	vRadius_.resize(MAX_LEVEL);
	vRadius_[0] = dBaseRadius_;
	for (int i = 1; i < MAX_LEVEL; ++i)
	{
		vRadius_[i] = vRadius_[i - 1] * 0.5f;
	}

	// get bounds
	vBoundsMin_.resize(uiDimension_);
	vBoundsMax_.resize(uiDimension_);

	for (int i = 0; i < uiDimTranslation; ++i)
	{
		vBoundsMin_[i] = -maxNeighborCellInRow * vCellWidth_[0] * 0.5;
		vBoundsMax_[i] =  maxNeighborCellInRow * vCellWidth_[0] * 0.5;
	}
	for (int i = uiDimTranslation; i < uiDimension_; ++i)
	{
		vBoundsMin_[i] = -1.0;
		vBoundsMax_[i] =  1.0;
	}

	vvDimensionSize_.resize(MAX_LEVEL);
	vvAccumulatedDimensionSize_.resize(MAX_LEVEL);
	for (int k = 0; k < MAX_LEVEL; ++k)
	{
		vvDimensionSize_[k].resize(uiDimension_);
		vvAccumulatedDimensionSize_[k].resize(uiDimension_);
		for (int i = 0; i < uiDimension_; ++i)
		{
			double dRange = vBoundsMax_[i] - vBoundsMin_[i];
			vvDimensionSize_[k][i] = (int)(ceil(dRange / vCellWidth_[k]));
			vvAccumulatedDimensionSize_[k][i] = vvDimensionSize_[k][i];

			for (int j = 0; j < i; ++j)
			{
				vvAccumulatedDimensionSize_[k][i] *= vvDimensionSize_[k][j];
			}
		}
	}
	uiNumTotalCells_ = vvAccumulatedDimensionSize_[0][uiDimension_ - 1];

	vBaseCells_.resize(vvAccumulatedDimensionSize_[0][uiDimension_ - 1]);
	vvAdditionalCells_.resize(MAX_LEVEL);

	// TODO: should be enough (or pointers will be invalid)
	vSamples_.reserve(100000);

	nnTemplate_.reset(new ompl::NearestNeighborsGNAT<const base::State*>());
	nnTemplate_->setDistanceFunction(boost::bind(&PoissonSampler::distanceFunction, this, _1, _2));

	initialized_ = true;
}

void PoissonSampler::createAdaptiveTemplate()
{
	clearSamples();

	vBaseCells_.resize(vvAccumulatedDimensionSize_[0][uiDimension_ - 1]);
	for (int j = 0; j < vvAccumulatedDimensionSize_[0][uiDimension_ - 1]; ++j)
	{
		vBaseCells_[j].baseId = j;
		vBaseCells_[j].level = 0;
		vBaseCells_[j].sampleId = -1;
		idToIndices(j, vBaseCells_[j].indices, 0);
	}
	
	std::vector<double> vCenter(si_->getStateSpace()->getSerializationLength() / sizeof(double));
	for (int i = 0; i < uiDimension_; ++i)
		vCenter[i] = 0.0;

	base::State* state = si_->allocState();
	si_->getStateSpace()->deserialize(state, &(vCenter[0]));
	si_->getStateSpace()->enforceBounds(state);

	int centerCellId = positionToId(vCenter, 0);
	GridCell* centerCell = getGridCell(centerCellId, 0);

	PoissonState* centerState = addState(centerCellId, 0, state, false);
	centerCell->sampleId = centerState->getId();
	nnTemplate_->add(centerState->getState());

	int failure = 0;
	int trials = 0;
	while (failure < 10000)
	{
		int randomCellId = rng_.uniformInt(0, uiNumTotalCells_-1);
		GridCell* cell = getGridCell(randomCellId, 0);


		if (cell && cell->GetStatus() == CS_NOSAMPLE)
		{
			randomPointInCell(*cell, state, vCellWidth_[0]);

			double distanceToCenter = distanceFunction(state, centerState->getState());
			if (isDiskFree(state, *cell, 1, true) && distanceToCenter > 0.5 * dBaseRadius_ && distanceToCenter < 1.5 * dBaseRadius_)
			{
				PoissonState* poissonState = addState(randomCellId, 0, state, false);
				cell->sampleId = poissonState->getId();
				nnTemplate_->add(poissonState->getState());
				//std::cout << "Sample added" << std::endl;
			}
			else
			{
				++failure;
			}
			++trials;
			//std::cout << failure << "/" << trials << " Failures" << std::endl;
		}
	}

	/*
	for (unsigned int i = 0; i < vvSamples_.size(); ++i)
		std::cout << vvSamples_[i].p[0] << ' ' << vvSamples_[i].p[1] << ' ' << vvSamples_[i].p[2] << std::endl;
		*/

	vvTemplateSamples_.reserve(vSamples_.size()-1);
	vvTemplateSamples_.insert(vvTemplateSamples_.end(), vSamples_.begin()+1, vSamples_.end());

	nnTemplate_->remove(centerState->getState());

	vSamples_.clear();
	delete centerState;

	vvAdditionalCells_.clear();
	vActiveCells_.clear();
	vActiveCells2_.clear();

	si_->freeState(state);
}

// utility functions
////////////////////////////////////////////////////////////////////////////////

bool PoissonSampler::idToIndices(int id, std::vector<int>& indices, int uiLevel) const
{
	if (!initialized_)
		return false;

	indices.resize(uiDimension_);

	indices[0] = id % vvDimensionSize_[uiLevel][0];
	for (unsigned int i = 1; i < uiDimension_; ++i)
	{
		indices[i] = (id / vvAccumulatedDimensionSize_[uiLevel][i - 1] ) % vvDimensionSize_[uiLevel][i];
	}

	return true;
}

int PoissonSampler::positionToId(const std::vector<double>& vPos, int uiLevel) const
{
	std::vector<int> vIndices(uiDimension_);
	for (int i = 0; i < uiDimension_; ++i)
	{
		vIndices[i] = (int)((vPos[i] - vBoundsMin_[i]) / vCellWidth_[uiLevel]);
	}
	int uiId;
	indicesToId(vIndices, uiId, uiLevel);

	return uiId;
}

bool PoissonSampler::indicesToId(const std::vector<int>& indices, int& id, int uiLevel) const
{
	if (!initialized_)
		return false;

	id = indices[0];
	for (int i = 1; i < uiDimension_; ++i)
	{
		id += indices[i] * vvAccumulatedDimensionSize_[uiLevel][i - 1];
	}

	if (id >= vvAccumulatedDimensionSize_[uiLevel][uiDimension_ - 1])
	{
		return false;
	}

	return true;
}

void PoissonSampler::clearSamples()
{
	for (int i = 0; i < vSamples_.size(); ++i)
	{
		delete vSamples_[i];
	}
	vSamples_.clear();
}

PoissonState* PoissonSampler::getNeighborSample(PoissonState* from, PoissonState* to)
{
	PoissonState* neighbor = NULL;

	double dFromToDist = distanceFunction(from->getState(), to->getState());
	
	base::State* intersect = si_->allocState();
	si_->getStateSpace()->interpolate(from->getState(), to->getState(), vRadius_[from->getLevel()]/dFromToDist, intersect);
	si_->getStateSpace()->enforceBounds(intersect);

	PoissonState poissonIntersectState(si_, intersect, -1, -1);

	std::vector<PoissonState*> neighbors;
	nn_->nearestR(&poissonIntersectState, vRadius_[0], neighbors);
	
	double dMinDistance = dFromToDist;
	for (unsigned int i = 0; i < neighbors.size(); ++i)
	{
		if (distanceFunction(neighbors[i]->getState(), intersect) > vRadius_[neighbors[i]->getLevel()])
			continue;

		double dDistance = distanceFunction(neighbors[i]->getState(), to->getState());
		if (dDistance < dMinDistance)
		{
			dMinDistance = dDistance;
			neighbor = neighbors[i];
		}
	}

	if (neighbor == NULL)
		neighbor = addState(-1, from->getLevel(), intersect, true);

	si_->freeState(intersect);

	if (neighbor && neighbor->isTreeNode())
		return NULL;

	return neighbor;
}

PoissonState* PoissonSampler::collisionDetected(PoissonState* from, PoissonState* to, base::State* free)
{
	if (subdivideDisk(to, free))
	{
		to->increaseLevel();
		//sampleManager_.ShrinkSample(to);
	}
	subdivideDisk(from, free);

	return NULL;
}

bool PoissonSampler::subdivideDisk(PoissonState* state, base::State* free)
{
	if (state->getLevel() == MAX_LEVEL - 1 || state == goalState_)
		return false;

	if (distanceFunction(state->getState(), free) < vRadius_[state->getLevel()])
	{
		applyAdaptiveTemplate(state, free);
	}

	return true;
}

void PoissonSampler::scaledDiff(const base::State* from, const base::State* to, base::State* diff, double scale) const
{
	unsigned int uiDimTranslation = uiDimension_;
	unsigned int uiDimRotation = 0;
	if (si_->getStateSpace()->isCompound())
	{
		uiDimTranslation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->getDimension();
		uiDimRotation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->getDimension();
	}

	// diff = (to - from) * scale

	// translation
	{
		std::vector<double> vDiff(uiDimTranslation);
		std::vector<double> vFrom(uiDimTranslation);
		std::vector<double> vTo(uiDimTranslation);

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->serialize(&(vFrom[0]), from->as<base::CompoundState>()->components[0]);
		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->serialize(&(vTo[0]), to->as<base::CompoundState>()->components[0]);

		for (int i = 0; i < uiDimTranslation; ++i)
			vDiff[i] = scale * (vTo[i] - vFrom[i]);

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->deserialize(diff->as<base::CompoundState>()->components[0], &(vDiff[0]));
	}

	// rotation
	{
		// SO3 needs 4 variables (quaternion)
		std::vector<double> vDiff(uiDimRotation + 1);
		std::vector<double> vFrom(uiDimRotation + 1);
		std::vector<double> vTo(uiDimRotation + 1);

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->serialize(&(vFrom[0]), from->as<base::CompoundState>()->components[1]);
		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->serialize(&(vTo[0]), to->as<base::CompoundState>()->components[1]);

		
		switch (uiDimRotation)
		{
		case 0:
			break;
		case 1:
			vDiff[0] = (vTo[0] - vFrom[0]);
			break;
		case 3:
		{
			// compute norm
			double nrmSqr = vFrom[0] * vFrom[0] + vFrom[1] * vFrom[1] + vFrom[2] * vFrom[2] + vFrom[3] * vFrom[3];
			double nrm = (fabs(nrmSqr - 1.0) > std::numeric_limits<double>::epsilon()) ? sqrt(nrmSqr) : 1.0;

			std::vector<double> inv(uiDimRotation + 1);

			// compute inverse of from
			inv[0] = -vFrom[0] / nrm;
			inv[1] = -vFrom[1] / nrm;
			inv[2] = -vFrom[2] / nrm;
			inv[3] = vFrom[3] / nrm;

			// from^-1 * to
			vDiff[0] = inv[3] * vTo[0] + inv[0] * vTo[3] + inv[1] * vTo[2] - inv[2] * vTo[1];
			vDiff[1] = inv[3] * vTo[1] + inv[1] * vTo[3] + inv[2] * vTo[0] - inv[0] * vTo[2];
			vDiff[2] = inv[3] * vTo[2] + inv[2] * vTo[3] + inv[0] * vTo[1] - inv[1] * vTo[0];
			vDiff[3] = inv[3] * vTo[3] - inv[0] * vTo[0] - inv[1] * vTo[1] - inv[2] * vTo[2];
		}
			break;
		default:
			assert(false);
			break;
		}

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->deserialize(diff->as<base::CompoundState>()->components[1], &(vDiff[0]));
	}
	si_->getStateSpace()->enforceBounds(diff);
}

void PoissonSampler::scaledAdd(const base::State* from, const base::State* add, base::State* sum, double scale) const
{
	unsigned int uiDimTranslation = uiDimension_;
	unsigned int uiDimRotation = 0;
	if (si_->getStateSpace()->isCompound())
	{
		uiDimTranslation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->getDimension();
		uiDimRotation = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->getDimension();
	}

	// sum = from + add * scale

	// translation
	{
		std::vector<double> vSum(uiDimTranslation);
		std::vector<double> vFrom(uiDimTranslation);
		std::vector<double> vAdd(uiDimTranslation);

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->serialize(&(vFrom[0]), from->as<base::CompoundState>()->components[0]);
		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->serialize(&(vAdd[0]), add->as<base::CompoundState>()->components[0]);

		for (int i = 0; i < uiDimTranslation; ++i)
			vSum[i] = vFrom[i] + scale * vAdd[i];

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0)->deserialize(sum->as<base::CompoundState>()->components[0], &(vSum[0]));
	}

	// rotation
	{
		// SO3 needs 4 variables (quaternion)
		std::vector<double> vSum(uiDimRotation + 1);
		std::vector<double> vFrom(uiDimRotation + 1);
		std::vector<double> vAdd(uiDimRotation + 1);

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->serialize(&(vFrom[0]), from->as<base::CompoundState>()->components[1]);
		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->serialize(&(vAdd[0]), add->as<base::CompoundState>()->components[1]);

		switch (uiDimRotation)
		{
		case 0:
			break;
		case 1:
			vSum[0] = vFrom[0] + vAdd[0];
			break;
		case 3:
			// from * add(from^-1 * to)
			vSum[0] = vFrom[3] * vAdd[0] + vFrom[0] * vAdd[3] + vFrom[1] * vAdd[2] - vFrom[2] * vAdd[1];
			vSum[1] = vFrom[3] * vAdd[1] + vFrom[1] * vAdd[3] + vFrom[2] * vAdd[0] - vFrom[0] * vAdd[2];
			vSum[2] = vFrom[3] * vAdd[2] + vFrom[2] * vAdd[3] + vFrom[0] * vAdd[1] - vFrom[1] * vAdd[0];
			vSum[3] = vFrom[3] * vAdd[3] - vFrom[0] * vAdd[0] - vFrom[1] * vAdd[1] - vFrom[2] * vAdd[2];
			break;
		default:
			assert(false);
			break;
		}

		si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1)->deserialize(sum->as<base::CompoundState>()->components[1], &(vSum[0]));
	}
	si_->getStateSpace()->enforceBounds(sum);
}

void PoissonSampler::applyAdaptiveTemplate(PoissonState* state, base::State* free)
{
	double scale = (double)(1 << state->getLevel());

	base::State* scaledDiffState = si_->allocState();
	scaledDiff(state->getState(), free, scaledDiffState, scale);

	if (const base::State* temp = nnTemplate_->nearest(scaledDiffState))
	{
		base::State* scaledAddState = si_->allocState();

		scaledAdd(state->getState(), temp, scaledAddState, 1.0 / scale);

		// goal position should be pickable
		PoissonState* poissonState = addState(-1, state->getLevel() + 1, scaledAddState, true);
		//sampleManager_.AddSample(temp->id, sample->level + 1);

		si_->freeState(scaledAddState);
	}

	si_->freeState(scaledDiffState);
}

void PoissonSampler::stateAccepted(PoissonState* state)
{
	state->setTreeNode(true);
}

}